'''
Author: your name
Date: 2022-02-24 12:35:14
LastEditTime: 2022-02-24 12:44:10
LastEditors: Please set LastEditors
Description: 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
FilePath: \other-algorithm-optimization\BPSO\BPSO.py
'''

from cmath import cos, exp, sin
import random
from turtle import end_fill
import numpy as np
import matplotlib.pyplot as plt
from globalVal import *


class BPSO(object):
    def __init__(self, Fobj, n_dim, lb, ub, max_iter=200, size_pop=100, c1=1.5, c2=1.5, Wmax=0.8, Wmin=0.4, Vmax = 10, Vmin = -10):
        '''
        func: 适应函数
        n_dim: 粒子维数
        max_iter: 最大迭代次数
        size_pop: 群体粒子个数
        c1: 个体学习因子
        c2: 群体学习因子
        Wmax: 惯性权重最大值
        Wmin: 惯性权重最小值
        lb: 位置最小值
        ub: 位置最大值
        Vmax: 速度最大值
        Vmin: 速度最小值
        '''
        self.Fobj = Fobj
        self.n_dim = n_dim
        self.max_iter = max_iter
        self.size_pop = size_pop
        self.c1 = c1
        self.c2 = c2
        self.Wmax = Wmax
        self.Wmin = Wmin
        self.ub = ub
        self.lb = lb
        self.Vmax = Vmax
        self.Vmin = Vmin
        self.x = None
        self.v = None
        self.p = None
        self.pbest = None
        self.g = None
        self.gbest = None
        self.gb = []

    def Initialization(self):
        self.x = np.random.randint(0, 2, (self.size_pop, self.n_dim))
        self.v = np.random.rand(self.size_pop, self.n_dim) * (self.Vmax - self.Vmin) + self.Vmin
        self.p = self.x
        self.pbest = np.ones(self.size_pop)
        for i in range(self.size_pop):
            stepFlag, self.pbest[i] = self.Fobj.step_online(self.x[i, :])
            if stepFlag != 2:
                self.pbest[i] = -1
        self.g = np.ones(self.n_dim)
        self.gbest = self.pbest[0]
        for i in range(self.size_pop):
            if self.pbest[i] < self.gbest:
                self.g = self.p[i, :]
                self.gbest = self.pbest[i]
        pass

    def run(self):
        self.Initialization()
        for i in range(self.max_iter):
            for j in range(self.size_pop):
                # 更新个体最优位置和最优值
                stepFlag, fitness = self.Fobj.step_online(self.x[j, :])
                if stepFlag != 2:
                    fitness = -1
                if fitness < self.pbest[j]:
                    self.p[j, :] = self.x[j, :]
                    self.pbest[j] = fitness
                # 更新全局最优位置和最优值
                if fitness < self.gbest:
                    self.g = self.p[j, :]
                    self.gbest = fitness
                # 计算动态惯性权重值
                w = self.Wmax - (self.Wmax - self.Wmin) * i / self.max_iter
                # 更新速度
                self.v[j, :] = w * self.v[j, :] + self.c1 * np.random.rand() *\
                        (self.p[j, :] - self.x[j, :]) + self.c2 * np.random.rand() * (self.g - self.x[j, :])
                # 边界条件处理
                # for ii in range(self.n_dim):
                #     if self.v[j, ii] < self.Vmin or self.v[j, ii] > self.Vmax:
                #         self.v[j, ii] = np.random.rand() * (self.Vmax, self.Vmin) + self.Vmin
                vx = 1. / (1 + np.exp(-self.v[j, :]))
                for jj in range(self.n_dim):
                    if vx[jj] > np.random.rand():
                        self.x[j, jj] = 1
                    else:
                        self.x[j, jj] = 0
            print(self.gbest)
            writer.add_scalar("BPSO/best fitness", -self.gbest, i)
            self.gb.append(self.gbest)
    
    def plot(self):
        plt.figure(figsize=(10, 8))
        plt.xlabel("迭代次数")
        plt.ylabel("适应度值")
        plt.title("适应度进化曲线")
        plt.plot(self.gb)
        plt.show()